#include "StdAfx.h"
#include "../System.h"
#include "../LoTray.h"

namespace lo {

	void TrayApp::PrepareStructure(UINT id, NOTIFYICONDATAW&	notifyIconData)
	{
		memset((void *)&notifyIconData, 0, sizeof(NOTIFYICONDATAW));
#if NTDDI_VERSION >= NTDDI_WINXP
		notifyIconData.cbSize = NOTIFYICONDATAW_V2_SIZE;
#else
		notifyIconData.cbSize = sizeof(NOTIFYICONDATA);
#endif
		notifyIconData.hWnd = hwnd;
		notifyIconData.uID = id;
		notifyIconData.uVersion = NOTIFYICON_VERSION;
	}
	void TrayApp::Info( UINT id, const char *title, const char *message, UINT milliseconds, DWORD flags )
	{
		NOTIFYICONDATAW	notifyIconData;
		PrepareStructure(id,notifyIconData);
		notifyIconData.uFlags = NIF_INFO;
		wcsncpy(notifyIconData.szInfo,wstring(message),_countof(notifyIconData.szInfo));
		wcsncpy(notifyIconData.szInfoTitle,wstring(title),_countof(notifyIconData.szInfoTitle));
		notifyIconData.uTimeout = milliseconds;
		notifyIconData.dwInfoFlags = flags;
		notifyIconData.uFlags = NIF_INFO;
		Shell_NotifyIconW(NIM_MODIFY,&notifyIconData);
	}
	void TrayApp::Tooltip( UINT id, const char *tip )
	{
		NOTIFYICONDATAW	notifyIconData;
		PrepareStructure(id,notifyIconData);
		notifyIconData.uFlags = NIF_TIP;
		wcsncpy(notifyIconData.szTip,wstring(tip),_countof(notifyIconData.szTip));
		Shell_NotifyIconW(NIM_MODIFY,&notifyIconData);
	}
	// TRAY APP FUNCTION
	TrayApp::TrayApp(HINSTANCE hinstance)
	{
		this->hinstance = hinstance;
		this->currentId = 0;
		CreateContainerWindow();
	}
	TrayApp::~TrayApp()
	{
		for( int i = this->currentId - 1; i >= 0; --i)
		{
			NOTIFYICONDATAW	notifyIconData;
			PrepareStructure((UINT)i, notifyIconData);
			Shell_NotifyIconW(NIM_DELETE,&notifyIconData);
			hwnd = NULL;
		}
	}
	UINT TrayApp::AddIcon(HICON hicon, HMENU leftMenu, HMENU rightMenu)
	{
		this->leftMenu[this->currentId] = leftMenu;
		this->rightMenu[this->currentId] = rightMenu;
		NOTIFYICONDATAW	notifyIconData;
		PrepareStructure(this->currentId,notifyIconData);
		notifyIconData.uFlags |= NIF_MESSAGE;
		notifyIconData.uCallbackMessage = WM_ICON_NOTIFY;
		if(hicon)
		{
			notifyIconData.uFlags |= NIF_ICON;
			notifyIconData.hIcon = hicon;
		}
		Shell_NotifyIconW(NIM_ADD,&notifyIconData);
		UINT id = this->currentId;
		this->currentId++;
		return id;
	}

	const	wchar_t TrayApp::windowClassBase[] = {L"TrayApp"};
	wchar_t	TrayApp::windowClass[100];
	int		TrayApp::windowClassLength;

	HWND TrayApp::CreateContainerWindow()
	{
		WCHAR	moduleFileName[MAX_PATH];
		DWORD	i;
		for( i = GetModuleFileNameW(NULL, moduleFileName, _countof(moduleFileName)) - 1; (i >= 0) && (moduleFileName[i-1]!='\\'); --i );
		windowClassLength = _snwprintf( windowClass, _countof(windowClass), L"%s-%s", windowClassBase, moduleFileName + i );

		// Register the window class
		WNDCLASSEXW wcex;

		wcex.cbSize = sizeof(WNDCLASSEX); 

		wcex.style			= CS_HREDRAW | CS_VREDRAW | CS_DBLCLKS;
		wcex.lpfnWndProc	= (WNDPROC)WindowProcedure;
		wcex.cbClsExtra		= 0;
		wcex.cbWndExtra		= 0;
		wcex.hInstance		= hinstance;
		wcex.hIcon			= 0;
		wcex.hCursor		= 0;
		wcex.hbrBackground	= 0;
		wcex.lpszMenuName	= 0;
		wcex.lpszClassName	= windowClass;
		wcex.hIconSm		= 0;

		RegisterClassExW(&wcex);

		hwnd = CreateWindowW(windowClass, L"", WS_POPUP,
			CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT, CW_USEDEFAULT, NULL, 0, hinstance, 0);
		// Save a pointer to the tray app.
		SetWindowLongPtr(hwnd,GWLP_USERDATA,(LONG)(LONG_PTR)this);

		return hwnd;
	}

	bool TrayApp::OnlyOne()
	{
		onlyOne = true;
		EnumWindows(OnlyOneEnumWindowsProc,reinterpret_cast<LPARAM>(this));
		if( !onlyOne )
		{
			MSG msg;
			// If timer message is received before kickback from remote, then kill the remote process
			DWORD	processId;
			GetWindowThreadProcessId(hwnd, &processId);
			SetTimer( hwnd, (UINT_PTR)processId, 2000, NULL );
			while (GetMessage(&msg, NULL, 0, 0))
			{
				TranslateMessage(&msg);
				DispatchMessage(&msg);
			}
			onlyOne = (msg.wParam != 0);
		}
		return onlyOne;
	}

	BOOL CALLBACK TrayApp::OnlyOneEnumWindowsProc( __in HWND hwnd, __in LPARAM lParam )
	{
		WCHAR	remoteClass[MAX_PATH];
		if( GetClassNameW(hwnd,remoteClass,_countof(windowClass)) == (windowClassLength))
		{
			TrayApp *tray = reinterpret_cast<TrayApp *>(lParam);
			if( (tray->hwnd != hwnd) && (wcscmp(remoteClass,windowClass) == 0) )
			{
				PostMessage( hwnd, WM_ISALIVE, (WPARAM)0, (WPARAM)tray->hwnd);
				tray->onlyOne = false;
			}
		}
		return TRUE;
	}

	int TrayApp::Run(Callback callback)
	{
		MSG msg;
		this->callback = callback;
		ShowWindow(hwnd, SW_HIDE);
		UpdateWindow(hwnd);
		while (GetMessage(&msg, NULL, 0, 0))
		{
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}

		return (int) msg.wParam;
	}

	void TrayApp::TrackMenu(HMENU hmenu)
	{
		POINT	point;
		GetCursorPos(&point);
		SetForegroundWindow(hwnd);
		TrackPopupMenu(GetSubMenu(hmenu,0), TPM_LEFTBUTTON, point.x, point.y, 0, hwnd, NULL);
		PostMessage(hwnd, WM_NULL, 0, 0);
	}

	LRESULT CALLBACK TrayApp::WindowProcedure(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
	{
		int wmId, wmEvent;
		PAINTSTRUCT ps;
		HDC hdc;
		TrayApp	*trayApp;

		switch (message)
		{
		case WM_ICON_NOTIFY:
			trayApp = reinterpret_cast<TrayApp *>(GetWindowLongPtr(hWnd,GWLP_USERDATA));
			{
				switch(lParam)
				{
				case WM_MOUSEMOVE:
					trayApp->callback(MouseHover,message,lParam);
					break;
				case WM_LBUTTONDOWN:
					if( trayApp->leftMenu[wParam] )
						trayApp->TrackMenu(trayApp->leftMenu[wParam]);
					else
						trayApp->callback(LeftMouseClick,message,lParam);
					break;
				case WM_RBUTTONDOWN:
					if( trayApp->rightMenu[wParam] )
						trayApp->TrackMenu(trayApp->rightMenu[wParam]);
					else
						trayApp->callback(RightMouseClick,message,lParam);
					break;
				}
			}
			break;
		case WM_ISALIVE:
			PostMessage((HWND)lParam, WM_QUITFROMREMOTE, 0, (LPARAM)hWnd);
			break;
		case WM_QUITFROMREMOTE:
			PostQuitMessage(0);	// We got a healthy kickback from remote app, so we should die
			SetActiveWindow((HWND)lParam);	// Set remote window as active.
			break;
		case WM_TIMER:
			HANDLE hHandle;
			if( hHandle = ::OpenProcess(PROCESS_ALL_ACCESS,0,(DWORD)wParam) )
			{
				::TerminateProcess(hHandle,1);
			}
			else
				MessageBoxA(hWnd,"Please kill the remote process.  It is dead.", "", MB_OK);
			PostQuitMessage(1);	// We quit here, because we need to go to the second loop, we are just in the test loop at this point.
			break;
		case WM_COMMAND:
			wmId    = LOWORD(wParam);
			wmEvent = HIWORD(wParam);
			trayApp = reinterpret_cast<TrayApp *>(GetWindowLongPtr(hWnd,GWLP_USERDATA));
			if( !trayApp->callback(Command, wmId, lParam) )
			{	// If we can't process the message, then let the default window proc handle it.
				return DefWindowProc(hWnd,message,wParam,lParam);
			}
			break;
		case WM_PAINT:
			hdc = BeginPaint(hWnd, &ps);
			EndPaint(hWnd, &ps);
			break;
		case WM_DESTROY:
			PostQuitMessage(0);
			break;
		default:
			return DefWindowProc(hWnd, message, wParam, lParam);
		}
		return 0;
	}
}
